"""
Action space implementation for the Claude Code POMDP.

This module implements the action space A as defined in the theoretical framework,
including tool actions, parameter validation, and concurrency specifications.
"""

using LinearAlgebra
using JSON3
using DataStructures

# ============================================================================
# Action Space Definition and Operations
# ============================================================================

"""
    ClaudeCodeActionSpace

Represents the complete action space for the Claude Code system.
Implements tool selection, parameterization, and concurrency management.
"""
struct ClaudeCodeActionSpace
    tools::Vector{Tool}
    tool_index::Dict{String, Int}  # tool name -> index
    compatibility_matrix::BitMatrix  # tool compatibility for concurrency
    parameter_validators::Dict{String, Function}  # tool -> validation function
    max_concurrent_tools::Int
    
    function ClaudeCodeActionSpace(tools::Vector{Tool}, max_concurrent::Int = 4)
        tool_index = Dict(tool.name => i for (i, tool) in enumerate(tools))
        compat_matrix = build_compatibility_matrix(tools)
        validators = build_parameter_validators(tools)
        
        new(tools, tool_index, compat_matrix, validators, max_concurrent)
    end
end

"""
    build_compatibility_matrix(tools::Vector{Tool}) -> BitMatrix

Build matrix indicating which tools can run concurrently.
"""
function build_compatibility_matrix(tools::Vector{Tool})
    n = length(tools)
    matrix = falses(n, n)
    
    for i in 1:n, j in 1:n
        if i == j
            matrix[i, j] = false  # Tool cannot run concurrently with itself
        else
            matrix[i, j] = are_tools_compatible(tools[i], tools[j])
        end
    end
    
    return matrix
end

"""
    are_tools_compatible(tool1::Tool, tool2::Tool) -> Bool

Determine if two tools can run concurrently safely.
"""
function are_tools_compatible(tool1::Tool, tool2::Tool)
    # Tools are incompatible if:
    # 1. Either is high risk
    # 2. Both modify the same resources
    # 3. One depends on the other's output
    
    if tool1.safety_level >= HIGH_RISK || tool2.safety_level >= HIGH_RISK
        return false
    end
    
    # File system conflicts
    if tool1.category == EDIT_TOOLS && tool2.category == EDIT_TOOLS
        return false  # Two edit tools might conflict
    end
    
    # Execution conflicts
    if tool1.category == EXEC_TOOLS && tool2.category == EXEC_TOOLS
        return false  # Multiple executions might interfere
    end
    
    # Check explicit dependencies
    if tool2.name in tool1.dependencies || tool1.name in tool2.dependencies
        return false
    end
    
    # Default to compatible if both are marked as concurrent_compatible
    return tool1.concurrent_compatible && tool2.concurrent_compatible
end

"""
    build_parameter_validators(tools::Vector{Tool}) -> Dict{String, Function}

Build parameter validation functions for each tool.
"""
function build_parameter_validators(tools::Vector{Tool})
    validators = Dict{String, Function}()
    
    for tool in tools
        validators[tool.name] = create_validator(tool)
    end
    
    return validators
end

"""
    create_validator(tool::Tool) -> Function

Create a parameter validation function for a specific tool.
"""
function create_validator(tool::Tool)
    schema = tool.parameter_schema
    
    return function validate_params(params::Dict{String, Any})
        errors = String[]
        
        # Check required parameters
        for (param_name, param_type) in schema
            if !haskey(params, param_name)
                push!(errors, "Missing required parameter: $param_name")
                continue
            end
            
            value = params[param_name]
            
            # Type checking
            if !isa(value, param_type)
                if param_type == String && !isa(value, AbstractString)
                    push!(errors, "Parameter $param_name must be a string, got $(typeof(value))")
                elseif param_type == Int && !isa(value, Integer)
                    push!(errors, "Parameter $param_name must be an integer, got $(typeof(value))")
                elseif param_type == Float64 && !isa(value, Real)
                    push!(errors, "Parameter $param_name must be a number, got $(typeof(value))")
                elseif param_type == Bool && !isa(value, Bool)
                    push!(errors, "Parameter $param_name must be a boolean, got $(typeof(value))")
                elseif param_type == Vector && !isa(value, AbstractVector)
                    push!(errors, "Parameter $param_name must be a vector, got $(typeof(value))")
                elseif param_type == Dict && !isa(value, AbstractDict)
                    push!(errors, "Parameter $param_name must be a dictionary, got $(typeof(value))")
                else
                    # Try to convert if reasonable
                    try
                        if param_type == String
                            params[param_name] = string(value)
                        elseif param_type == Float64
                            params[param_name] = Float64(value)
                        elseif param_type == Int
                            params[param_name] = Int(value)
                        end
                    catch
                        push!(errors, "Cannot convert parameter $param_name to $(param_type)")
                    end
                end
            end
        end
        
        # Tool-specific validation
        append!(errors, validate_tool_specific(tool, params))
        
        return isempty(errors), errors
    end
end

"""
    validate_tool_specific(tool::Tool, params::Dict{String, Any}) -> Vector{String}

Perform tool-specific parameter validation.
"""
function validate_tool_specific(tool::Tool, params::Dict{String, Any})
    errors = String[]
    
    if tool.category == FILE_TOOLS
        # Validate file paths
        for (key, value) in params
            if occursin("path", lowercase(key)) || occursin("file", lowercase(key))
                if isa(value, AbstractString)
                    if tool.name in ["write_file", "edit_file"] && !is_safe_path(value)
                        push!(errors, "Potentially unsafe file path: $value")
                    end
                end
            end
        end
        
    elseif tool.category == EXEC_TOOLS
        # Validate execution parameters
        if haskey(params, "command")
            command = params["command"]
            if isa(command, AbstractString)
                if contains_dangerous_patterns(command)
                    push!(errors, "Command contains potentially dangerous patterns")
                end
            end
        end
        
        if haskey(params, "timeout")
            timeout = params["timeout"]
            if isa(timeout, Real) && (timeout <= 0 || timeout > 3600)
                push!(errors, "Timeout must be between 0 and 3600 seconds")
            end
        end
        
    elseif tool.category == SEARCH_TOOLS
        # Validate search parameters
        if haskey(params, "pattern")
            pattern = params["pattern"]
            if isa(pattern, AbstractString)
                try
                    Regex(pattern)
                catch
                    push!(errors, "Invalid regular expression pattern")
                end
            end
        end
        
        if haskey(params, "max_results")
            max_results = params["max_results"]
            if isa(max_results, Integer) && (max_results <= 0 || max_results > 10000)
                push!(errors, "max_results must be between 1 and 10000")
            end
        end
    end
    
    return errors
end

"""
    is_safe_path(path::String) -> Bool

Check if a file path is safe for operations.
"""
function is_safe_path(path::String)
    # Disallow dangerous paths
    dangerous_patterns = [
        r"^\s*\/",          # Absolute paths starting with /
        r"\.\.\s*\/",       # Path traversal attempts
        r"\/etc\/",         # System configuration
        r"\/bin\/",         # System binaries
        r"\/usr\/bin\/",    # User binaries
        r"\/dev\/",         # Device files
        r"\/proc\/",        # Process information
        r"\/sys\/",         # System information
        r"^\s*~\/\.\w",     # Hidden files in home directory
    ]
    
    for pattern in dangerous_patterns
        if occursin(pattern, path)
            return false
        end
    end
    
    return true
end

"""
    contains_dangerous_patterns(command::String) -> Bool

Check if a command contains potentially dangerous patterns.
"""
function contains_dangerous_patterns(command::String)
    dangerous_patterns = [
        r"rm\s+-rf",           # Recursive deletion
        r"sudo\s+",            # Privilege escalation
        r"chmod\s+777",        # Dangerous permissions
        r">/dev/",             # Device access
        r"curl.*\|.*sh",       # Download and execute
        r"wget.*\|.*sh",       # Download and execute
        r"eval\s*\(",          # Code evaluation
        r"exec\s*\(",          # Code execution
        r"system\s*\(",        # System calls
        r"`.*`",               # Command substitution
        r"\$\(.*\)",           # Command substitution
        r"&\s*$",              # Background execution
        r";\s*$",              # Command chaining
        r"\|\s*$",             # Pipe chaining
    ]
    
    for pattern in dangerous_patterns
        if occursin(pattern, command)
            return true
        end
    end
    
    return false
end

"""
    create_action(action_space::ClaudeCodeActionSpace, tool_name::String, 
                  params::Dict{String, Any}, priority::Float64 = 1.0, 
                  timeout::Float64 = 60.0) -> Union{ToolAction, Nothing}

Create a validated tool action.
"""
function create_action(action_space::ClaudeCodeActionSpace, tool_name::String,
                      params::Dict{String, Any}, priority::Float64 = 1.0,
                      timeout::Float64 = 60.0)
    
    # Find the tool
    if !haskey(action_space.tool_index, tool_name)
        @error "Unknown tool: $tool_name"
        return nothing
    end
    
    tool_idx = action_space.tool_index[tool_name]
    tool = action_space.tools[tool_idx]
    
    # Validate parameters
    validator = action_space.parameter_validators[tool_name]
    is_valid, errors = validator(params)
    
    if !is_valid
        @error "Parameter validation failed for $tool_name" errors
        return nothing
    end
    
    # Create concurrency specification (initially no concurrency)
    concurrency_spec = falses(length(action_space.tools))
    
    return ToolAction(tool, params, concurrency_spec, priority, timeout)
end

"""
    create_concurrent_action(action_space::ClaudeCodeActionSpace, 
                           primary_action::ToolAction,
                           concurrent_tools::Vector{String}) -> Union{ToolAction, Nothing}

Create an action with specified concurrent tools.
"""
function create_concurrent_action(action_space::ClaudeCodeActionSpace,
                                 primary_action::ToolAction,
                                 concurrent_tools::Vector{String})
    
    if length(concurrent_tools) + 1 > action_space.max_concurrent_tools
        @error "Too many concurrent tools requested"
        return nothing
    end
    
    primary_idx = action_space.tool_index[primary_action.tool.name]
    concurrency_spec = copy(primary_action.concurrency_spec)
    
    # Check compatibility and set concurrency bits
    for tool_name in concurrent_tools
        if !haskey(action_space.tool_index, tool_name)
            @error "Unknown concurrent tool: $tool_name"
            return nothing
        end
        
        tool_idx = action_space.tool_index[tool_name]
        
        # Check if tools are compatible
        if !action_space.compatibility_matrix[primary_idx, tool_idx]
            @error "Tools $primary_action.tool.name and $tool_name are not compatible for concurrent execution"
            return nothing
        end
        
        concurrency_spec[tool_idx] = true
    end
    
    return ToolAction(primary_action.tool, primary_action.parameters, 
                     concurrency_spec, primary_action.priority, primary_action.timeout)
end

"""
    get_valid_concurrent_tools(action_space::ClaudeCodeActionSpace, 
                              tool_name::String) -> Vector{String}

Get list of tools that can run concurrently with the given tool.
"""
function get_valid_concurrent_tools(action_space::ClaudeCodeActionSpace, tool_name::String)
    if !haskey(action_space.tool_index, tool_name)
        return String[]
    end
    
    tool_idx = action_space.tool_index[tool_name]
    compatible_tools = String[]
    
    for (other_name, other_idx) in action_space.tool_index
        if other_name != tool_name && action_space.compatibility_matrix[tool_idx, other_idx]
            push!(compatible_tools, other_name)
        end
    end
    
    return compatible_tools
end

"""
    estimate_action_cost(action::ToolAction) -> Float64

Estimate the computational cost of executing an action.
"""
function estimate_action_cost(action::ToolAction)
    base_cost = action.tool.execution_cost
    
    # Adjust based on parameters
    param_complexity = 1.0
    for (key, value) in action.parameters
        if isa(value, AbstractString) && length(value) > 1000
            param_complexity *= 1.2  # Large string parameters
        elseif isa(value, AbstractVector) && length(value) > 100
            param_complexity *= 1.1  # Large array parameters
        elseif isa(value, AbstractDict) && length(value) > 50
            param_complexity *= 1.15  # Large dict parameters
        end
    end
    
    # Adjust based on safety level (higher risk = higher cost due to safety checks)
    safety_multiplier = if action.tool.safety_level == DANGEROUS
        3.0
    elseif action.tool.safety_level == HIGH_RISK
        2.0
    elseif action.tool.safety_level == MEDIUM_RISK
        1.5
    else
        1.0
    end
    
    return base_cost * param_complexity * safety_multiplier
end

"""
    estimate_action_latency(action::ToolAction) -> Float64

Estimate the execution latency of an action.
"""
function estimate_action_latency(action::ToolAction)
    base_latency = action.tool.average_latency
    
    # Adjust based on parameters
    if haskey(action.parameters, "timeout")
        timeout = action.parameters["timeout"]
        if isa(timeout, Real)
            base_latency = min(base_latency, Float64(timeout))
        end
    end
    
    # Adjust based on tool category
    category_multiplier = if action.tool.category == EXEC_TOOLS
        2.0  # Execution tools are generally slower
    elseif action.tool.category == SEARCH_TOOLS
        1.5  # Search can be slow on large codebases
    elseif action.tool.category == BUILD_TOOLS
        3.0  # Build operations are typically slow
    else
        1.0
    end
    
    return base_latency * category_multiplier
end

"""
    action_risk_level(action::ToolAction) -> Float64

Assess the risk level of an action on a scale [0, 1].
"""
function action_risk_level(action::ToolAction)
    base_risk = if action.tool.safety_level == DANGEROUS
        0.9
    elseif action.tool.safety_level == HIGH_RISK
        0.7
    elseif action.tool.safety_level == MEDIUM_RISK
        0.4
    elseif action.tool.safety_level == LOW_RISK
        0.1
    else
        0.05
    end
    
    # Increase risk for certain parameter patterns
    risk_multiplier = 1.0
    
    if action.tool.category == EXEC_TOOLS
        if haskey(action.parameters, "command")
            command = action.parameters["command"]
            if isa(command, AbstractString)
                if contains(command, "rm")
                    risk_multiplier *= 2.0
                end
                if contains(command, "sudo")
                    risk_multiplier *= 3.0
                end
            end
        end
    end
    
    if action.tool.category == EDIT_TOOLS
        if haskey(action.parameters, "file_path")
            path = action.parameters["file_path"]
            if isa(path, AbstractString) && !is_safe_path(path)
                risk_multiplier *= 2.0
            end
        end
    end
    
    return min(1.0, base_risk * risk_multiplier)
end

"""
    serialize_action(action::ToolAction) -> Dict{String, Any}

Serialize an action to a dictionary for storage/transmission.
"""
function serialize_action(action::ToolAction)
    return Dict{String, Any}(
        "tool_name" => action.tool.name,
        "parameters" => action.parameters,
        "concurrency_spec" => collect(action.concurrency_spec),
        "priority" => action.priority,
        "timeout" => action.timeout,
        "estimated_cost" => estimate_action_cost(action),
        "estimated_latency" => estimate_action_latency(action),
        "risk_level" => action_risk_level(action)
    )
end

"""
    deserialize_action(action_space::ClaudeCodeActionSpace, 
                      data::Dict{String, Any}) -> Union{ToolAction, Nothing}

Deserialize an action from a dictionary.
"""
function deserialize_action(action_space::ClaudeCodeActionSpace, data::Dict{String, Any})
    tool_name = data["tool_name"]
    parameters = data["parameters"]
    priority = get(data, "priority", 1.0)
    timeout = get(data, "timeout", 60.0)
    
    # Create base action
    action = create_action(action_space, tool_name, parameters, priority, timeout)
    if action === nothing
        return nothing
    end
    
    # Set concurrency specification if provided
    if haskey(data, "concurrency_spec")
        concurrency_bits = data["concurrency_spec"]
        if isa(concurrency_bits, AbstractVector) && length(concurrency_bits) == length(action_space.tools)
            action = ToolAction(action.tool, action.parameters, 
                              BitVector(concurrency_bits), action.priority, action.timeout)
        end
    end
    
    return action
end